use futures::stream::Stream; // 引入 Stream trait 实现
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::time::{Duration, Sleep};

/// 一个每次返回数字并异步延迟 1 秒的 Stream 示例。
pub struct MyStream {
    /// 当前已生成的元素计数，每次 poll_next 成功后加一
    count: u32,
    /// 流最多能生成多少个元素，到达后流结束
    max: u32,
    /// 用于异步延迟的定时器，每次生成新元素前都要等待 1 秒
    delay: Option<Pin<Box<Sleep>>>, // 用于异步延迟的定时器，每次生成新元素前都要等待 1 秒
}

impl MyStream {
    /// 创建一个新的 MyStream，max 指定最多返回多少个元素
    pub fn new(max: u32) -> Self {
        MyStream { count: 0, max, delay: None }
    }
}

impl Stream for MyStream {
    type Item = u32;

    /// poll_next 是 Stream 的核心方法，每次 poll 检查是否可以产生新元素。
    /// 如果延迟未完成，则返回 Pending，延迟完成后唤醒任务并返回新数据。
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        // 如果已达到最大计数，返回 None 表示流结束
        if this.count >= this.max {
            return Poll::Ready(None);
        }
        // 如果没有延迟，则创建一个新的异步延迟定时器
        if this.delay.is_none() {
            this.delay = Some(Box::pin(tokio::time::sleep(Duration::from_secs(1))));
        }
        // 检查延迟是否完成
        let delay = this.delay.as_mut().unwrap();
        match delay.as_mut().poll(cx) {
            // 延迟完成，返回新数据
            Poll::Ready(_) => {
                this.count += 1;
                this.delay = None; // 重置延迟，准备下一个元素
                Poll::Ready(Some(this.count))
            },
            // 延迟未完成，返回 Pending，等待唤醒
            Poll::Pending => Poll::Pending,
        }
    }
}

// 测试 stream
#[cfg(test)]
mod tests {
    use super::*;
    use futures::stream::StreamExt; // 仅在测试中引入 StreamExt trait 以使用 next() 方法

    #[tokio::test]
    async fn stream_test() {
        let mut my_stream = MyStream::new(5);
        // 异步迭代流，每次收到新数据都会打印
        while let Some(value) = my_stream.next().await {
            println!("Received: {}", value);
        }
        println!("Stream completed");
    }
}
